package CPU.rv64_1stage

import chisel3._
import chisel3.util._

object SignExt {
  def apply(a: UInt, len: Int) = {
    val aLen = a.getWidth
    val signBit = a(aLen-1)
    if (aLen >= len) a(len-1,0) else Cat(Fill(len - aLen, signBit), a)
  }
}

object ZeroExt {
  def apply(a: UInt, len: Int) = {
    val aLen = a.getWidth
    if (aLen >= len) a(len-1,0) else Cat(0.U((len - aLen).W), a)
  }
}

class Execution extends Module{
  val io = IO(new Bundle{
    val opcode = Input(UInt(7.W))
    val in1    = Input(UInt(64.W))
    val in2    = Input(UInt(64.W))
    val fuType = Input(UInt(3.W))
    val rs2    = Input(UInt(64.W))
    val out    = Output(UInt(64.W))
    val branhit = Output(Bool())
    val dmem   = new RamIO
  })
  val alu_out = MuxCase(0.U,Array(
    (io.opcode===ALUOpType.add)->(io.in1 + io.in2) ,
    (io.opcode===ALUOpType.sub)->(io.in1 - io.in2) ,
    (io.opcode===ALUOpType.and)->(io.in1 & io.in2) ,
    (io.opcode===ALUOpType.or )->(io.in1 | io.in2) ,
    (io.opcode===ALUOpType.xor)->(io.in1 ^ io.in2) ,
    (io.opcode===ALUOpType.slt)->(io.in1.asSInt() < io.in2.asSInt()).asUInt() ,
    (io.opcode===ALUOpType.sltu)->(io.in1 < io.in2).asUInt() ,
    (io.opcode===ALUOpType.sll)->(io.in1 << io.in2(5,0)),
    (io.opcode===ALUOpType.sra)->(io.in1.asSInt() >> io.in2(5,0)).asUInt() ,
    (io.opcode===ALUOpType.srl)->(io.in1 >> io.in2(5,0)),

    (io.opcode===ALUOpType.addw)->SignExt((io.in1 + io.in2)(31,0),64) ,
    (io.opcode===ALUOpType.subw)->SignExt((io.in1 - io.in2)(31,0),64) ,
    (io.opcode===ALUOpType.sllw)->SignExt((io.in1 << io.in2(4,0))(31,0),64) ,
    (io.opcode===ALUOpType.srlw)->SignExt((io.in1(31,0) >> io.in2(4,0)).asUInt(),64) ,
    (io.opcode===ALUOpType.sraw)->SignExt((io.in1(31,0).asSInt() >> io.in2(4,0)).asUInt(),64),
    (io.opcode===ALUOpType.jal) -> (io.in1 + io.in2),
    (io.opcode===ALUOpType.jalr)-> (io.in1 + io.in2),
  ))
  io.branhit := MuxCase(false.B,Array(
    (io.opcode===ALUOpType.beq)-> (io.in1 === io.in2),
    (io.opcode===ALUOpType.bne)-> (io.in1 =/= io.in2),
    (io.opcode===ALUOpType.blt)-> (io.in1.asSInt() < io.in2.asSInt()),
    (io.opcode===ALUOpType.bge)-> (io.in1.asSInt() >=io.in2.asSInt()),
    (io.opcode===ALUOpType.bltu)->(io.in1 < io.in2),
    (io.opcode===ALUOpType.bgeu)->(io.in1 >=io.in2),
    (io.opcode===ALUOpType.jal) -> true.B,
    (io.opcode===ALUOpType.jalr)-> true.B,
  ))
  val addroffset = WireInit(0.U(3.W))
  val offsetdata = WireInit(0.U(64.W))
//  offsetdata := MuxCase(io.dmem.rdata,Array(
//    (addroffset===)
//  ))
  offsetdata := io.dmem.rdata >> (addroffset<<3)
  val load_data = MuxCase(0.U,Array(
    (io.opcode===LSUOpType.lb)  -> SignExt(offsetdata(7,0),64),
    (io.opcode===LSUOpType.lbu) -> ZeroExt(offsetdata(7,0),64),
    (io.opcode===LSUOpType.lh)  -> SignExt(offsetdata(15,0),64),
    (io.opcode===LSUOpType.lhu) -> ZeroExt(offsetdata(15,0),64),
    (io.opcode===LSUOpType.lw)  -> SignExt(offsetdata(31,0),64),
    (io.opcode===LSUOpType.lwu) -> ZeroExt(offsetdata(31,0),64),
    (io.opcode===LSUOpType.ld)  -> offsetdata,
  ))

  io.out := Mux(io.fuType===FuType.lsu,load_data,alu_out)

  io.dmem.en   := io.fuType === FuType.lsu
  io.dmem.addr := io.in1 + io.in2
  addroffset   := io.dmem.addr(2,0)
  io.dmem.wen  := (io.fuType===FuType.lsu) && LSUOpType.isStore(io.opcode)
  io.dmem.wdata:= io.rs2 << (addroffset<<3)

  val mask = MuxCase(0.U,Array(
    (io.opcode===LSUOpType.sb) -> ZeroExt("hff".U,64),
    (io.opcode===LSUOpType.sh) -> ZeroExt("hffff".U,64),
    (io.opcode===LSUOpType.sw) -> ZeroExt("hffffffff".U,64),
    (io.opcode===LSUOpType.sd) -> (-1).S(64.W).asUInt(),
  ))
  io.dmem.wmask := mask << (addroffset<<3)

}
